package com.ruoyi.mp.handler;

import com.ruoyi.ai.factory.ModelChatService;
import com.ruoyi.ai.factory.ModelChatStrategyFactory;
import com.ruoyi.mp.handler.bo.TextBuilder;
import lombok.extern.slf4j.Slf4j;
import me.chanjar.weixin.common.session.WxSessionManager;
import me.chanjar.weixin.mp.api.WxMpService;
import me.chanjar.weixin.mp.bean.message.WxMpXmlMessage;
import me.chanjar.weixin.mp.bean.message.WxMpXmlOutMessage;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 简单实现公众号AI回复功能，
 * 先存储在本地，后面再考虑使用redis和落库
 */
@Slf4j
@Component
public class MsgHandler extends AbstractHandler {

    private final ConcurrentHashMap<String, String> responseCache = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<String, AtomicInteger> requestCountCache = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<String, Boolean> userStateCache = new ConcurrentHashMap<>(); // 新增状态缓存
    private final static int REQ_LIMIT = 20;
    private final static String HELP_MSG = "你好啊，旅行者，欢迎咨询，请输入以下指令关键词：" +
        "\n【请问】：咨询问题，例如请问提瓦特在哪里？" +
        "\n【1】：查看回复" +
        "\n【2】：查看今日剩余提问次数";
    private final static String DEFAULT_RES = "小派蒙还在思考中，再等等吧~";

    @Override
    public WxMpXmlOutMessage handle(WxMpXmlMessage wxMessage,
                                    Map<String, Object> context, WxMpService weixinService,
                                    WxSessionManager sessionManager) {
        String userMsg = wxMessage.getContent();
        String fromUser = wxMessage.getFromUser();

        if (isTriggered(userMsg)) {
            if (checkRequestLimit(fromUser)) {
                if (!canAskNewQuestion(fromUser)) {
                    return new TextBuilder().build("请先查看上一个问题的回复再继续提问", wxMessage, weixinService);
                }
                responseCache.put(fromUser, DEFAULT_RES);
                // 异步调用接口并存储回复结果
                CompletableFuture.runAsync(() -> {
                    ModelChatService modelChatService = ModelChatStrategyFactory.getEnableMfg();
                    String finalContent = userMsg + "(字数请控制在1024内)";
                    String res = modelChatService.chatStreamOnce(finalContent);
                    log.info("AI返回：{}", res);
                    if (res.length() > 1024) {
                        res = res.substring(0, 1024); // 截断超过1024个字符的部分
                    }
                    responseCache.put(fromUser, res);
                });
                setUserState(fromUser, false); // 设置为等待回复状态
                return new TextBuilder().build("小派蒙正在思考中...(大约30s)", wxMessage, weixinService);
            } else {
                return new TextBuilder().build("您今天已经使用了" + REQ_LIMIT + "次AI聊天功能，请明天再试", wxMessage, weixinService);
            }
        } else if (isGetRes(userMsg)) {
            // 获取回复结果
            String response = null;
            if (responseCache.containsKey(fromUser)) {
                response = responseCache.get(fromUser);
                if (!response.equals(DEFAULT_RES)) {
                    responseCache.remove(fromUser); // 清除已获取的回复结果
                    setUserState(fromUser, true); // 设置为可以提问状态
                }
            } else {
                response = "请先输入【请问】关键词提问哦~";
            }

            return new TextBuilder().build(response, wxMessage, weixinService);
        } else if (isCheckRemainingRequests(userMsg)) {
            // 查询剩余次数
            int remainingRequests = REQ_LIMIT - getRequestCount(fromUser);
            return new TextBuilder().build("您今天剩余 " + remainingRequests + " 次AI聊天功能", wxMessage, weixinService);
        }
        return new TextBuilder().build(HELP_MSG, wxMessage, weixinService);
    }

    /**
     * 开启AI聊天关键词
     */
    private boolean isTriggered(String message) {
        return message.startsWith("请问");
    }

    /**
     * 获取聊天回复关键词
     */
    private boolean isGetRes(String message) {
        return message.equals("1");
    }

    /**
     * 查询剩余次数关键词
     */
    private boolean isCheckRemainingRequests(String message) {
        return message.equals("2");
    }

    /**
     * 检查用户的请求次数是否超过限制
     */
    private boolean checkRequestLimit(String fromUser) {
        AtomicInteger count = requestCountCache.computeIfAbsent(fromUser, k -> new AtomicInteger(0));
        int currentCount = count.incrementAndGet();
        if (currentCount > REQ_LIMIT) {
            return false;
        }
        return true;
    }

    /**
     * 获取用户的请求次数
     */
    private int getRequestCount(String fromUser) {
        AtomicInteger count = requestCountCache.get(fromUser);
        return count != null ? count.get() : 0;
    }


    // 每天重置请求次数
    public void resetRequestCounts() {
        requestCountCache.clear();
    }

    /**
     * 设置用户状态
     */
    private void setUserState(String fromUser, boolean canAsk) {
        userStateCache.put(fromUser, canAsk);
    }

    /**
     * 检查用户是否可以提问
     */
    private boolean canAskNewQuestion(String fromUser) {
        Boolean state = userStateCache.getOrDefault(fromUser, true);
        return state;
    }
}
